"""
Author: xubing
Date: 2024-02-02 08:12:02
LastEditors: xubing
LastEditTime: 2024-03-13 07:39:09
Description: 
向量数据库相关操作
1. load本地向量模型
2. 向量数据库中collection构建
    a. naive build: standard method
    b. advance build: parent document retrieval(切分小(子)文档用来检索. 使用大(父)文档用来生成.)
"""

from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import InMemoryStore
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores import Milvus
from loguru import logger
from pymilvus import connections, utility
from pymilvus.exceptions import MilvusException

from core.text_splitter import ChineseRecursiveTextSplitter

fmt = "=== {:30} ==="

vector_store_config = {
    "parent_chunk_size": 1000,
    "child_chunk_size": 300,
}


class MilvusStore:
    def __init__(self, embedding_model_name_or_path, milvus_conn_args):
        # 连接向量数据库
        logger.info(fmt.format("Start connecting to Milvus..."))
        try:
            connections.connect(**milvus_conn_args)
            logger.info(fmt.format("Connect Success!"))
        except MilvusException as e:
            logger.error("Milvus连接出错,请检查连接.")
        self.embedding_function = self._get_hf_bge_embedding(
            embedding_model_name_or_path
        )
        self.milvus_conn_args = milvus_conn_args

    def _get_hf_bge_embedding(self, embedding_model_name_or_path):
        logger.info("|-> 使用Local Weight Embedding模型")
        if embedding_model_name_or_path is None:
            raise EnvironmentError
        model_kwargs = {"device": "cuda"}
        # set True to compute cosine similarity
        encode_kwargs = {"normalize_embeddings": True}
        embedding_function = HuggingFaceBgeEmbeddings(
            model_name=embedding_model_name_or_path,
            model_kwargs=model_kwargs,
            encode_kwargs=encode_kwargs,
            query_instruction="Represent this sentence for searching relevant passages: finance, investments, economics, real estate",
        )
        return embedding_function

    def _check(self, collection_name):
        # check if collection is exists
        has = utility.has_collection(collection_name)
        logger.info(
            "Check collection {} exist in Milvus: {}".format(collection_name, has)
        )
        return has

    def _save(self, vector_store):
        # 癔想
        vector_store.persist("./vector_store")

    def advance_build(self, collection_name, orig_docs, save=False):
        parent_chunksize = 1000
        child_chunksize = 300

        logger.info(f"开始新建{collection_name}")
        # 创建主文档分割器
        parent_splitter = ChineseRecursiveTextSplitter(
            chunk_size=parent_chunksize,
            chunk_overlap=0,
            separators=["\n\n", "\n", "(?<=\. )", " ", ""],
        )

        # 创建子文档分割器
        child_splitter = ChineseRecursiveTextSplitter(
            chunk_size=child_chunksize,
            chunk_overlap=0,
            separators=["\n\n", "\n", "(?<=\. )", " ", ""],
        )

        # 创建向量数据库对象
        vectorstore = Milvus(
            embedding_function=self.embedding_function,
            collection_name=collection_name,
            connection_args=self.milvus_conn_args,
            auto_id=True,
        )
        # 创建内存存储对象
        store = InMemoryStore()
        # 创建父文档检索器
        retriever = ParentDocumentRetriever(
            vectorstore=vectorstore,
            docstore=store,
            child_splitter=child_splitter,
            parent_splitter=parent_splitter,
            search_kwargs={"k": 4},
        )

        # 添加文档集
        retriever.add_documents(orig_docs)
        logger.info("Advance build Done!")
        # if save:
        #     vector_store = Milvus(
        #         collection_name=collection_name,  # 注意加载与创建这个参数不一样
        #         embedding_function=self.embedding_function,
        #         connection_args=self.milvus_conn_args,
        #     )
        #     self._save(vector_store)
        return retriever

    def get_milvusstore(self, collection_name, from_local=False):
        logger.info(f"开始加载{collection_name}")
        # if from_local:
        #     vector_store = Milvus(
        #         persist_directory="./vector_store",
        #         embedding_function=self.embedding_function,
        #     )
        #     return vector_store

        if self._check(collection_name):
            vectorstore = Milvus(
                collection_name=collection_name,  # 注意加载与创建这个参数不一样
                embedding_function=self.embedding_function,
                connection_args=self.milvus_conn_args,
            )
            return vectorstore

        else:
            logger.error("当前知识库未建立,请检查!")
            return

    def get_retriever(self, collection_name):
        vectorstore = self.get_milvusstore(collection_name)
        retriver = vectorstore.as_retriever()
        return retriver


if __name__ == "__main__":
    import os

    from fin_preprocess import DocPreprocessing

    file_path = "data/new"
    file_list = [os.path.join(file_path, x) for x in os.listdir(file_path)]
    docs = DocPreprocessing().batch_process(file_list)

    import yaml

    with open("config.yaml", "r") as f:
        config = yaml.safe_load(f)
    milvus_store = MilvusStore(
        embedding_model_name_or_path=config.get("embedding_model_name_or_path"),
        milvus_conn_args=config.get("milvus_conn_args"),
    )
    collection_name = "fin_emb"

    # retriever = milvus_store.advance_build(collection_name, docs, save=True)
    # retriever = milvus_store.get_retriever(collection_name)
    retriever = milvus_store.get_milvusstore(collection_name, from_local=True)
